#!/usr/bin/env python
# -*- coding: utf-8 -*-
# websocket 接口相关代码

from typing import List, Dict, Optional
from fastapi import APIRouter, WebSocket, Depends, WebSocketDisconnect, Cookie, Query, status


"""
websocket 链接管理
"""
class ConnectionManager:

    def __init__(self):
        """存放链接"""
        self.active_connections: List[Dict[str, WebSocket]] = []

    async def connect(self, user: str, ws: WebSocket):
        """链接"""
        await ws.accept()
        self.active_connections.append({"user": user, "ws": ws})

    async def disconnect(self, user: str, ws: WebSocket):
        """断开链接，移除"""
        self.active_connections.remove({"user": user, "ws": ws})

    @staticmethod
    async def send_personal_message(message: str, ws: WebSocket):
        """发送所有人消息"""
        await ws.send_text(message)


    async def send_other_message(self, message: dict, user: str):
        """发送个人消息"""
        for conn in self.active_connections:
            if conn["user"] == user:
                await conn["ws"].send_json(message)


    async def broadcast(self, data: str):
        """广播"""
        for conn in self.active_connections:
            await conn["ws"].send_text(data)


    async def broadcast_other(self, message: str, user: str):
        """广播(排除个人)"""
        for conn in self.active_connections:
            if conn["user"] != user:
                await conn["ws"].send_text(message)


    async def broadcast_json(self, data: dict):
        """广播json数据"""
        for conn in self.active_connections:
            await conn["ws"].send_json(data)

    @staticmethod
    async def send_personal_message_json(data: dict, ws: WebSocket):
        """发送所有人消息(json数据格式)"""
        await ws.send_json(data)


ws_manager = ConnectionManager()

websocket_router = APIRouter()



# 没对token做实际验证
async def get_cookie_or_token(
        websocket: WebSocket,
        session: Optional[str] = Cookie(None),
        token: Optional[str] = Query(None),
):
    if session is None and token is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
    return session or token

# 没对token做实际验证
@websocket_router.websocket("/ws/{user}")
async def websocket_json_data(websocket: WebSocket, user: str, cookie_or_token: str = Depends(get_cookie_or_token)):
    await ws_manager.connect(user, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            send_user = data["username"]
            if send_user:
                # await ws_manager.send_other_message(data, send_user)
                data["username"] = '你'
                await ws_manager.send_personal_message_json(data, websocket)
                data["username"] = user
                await ws_manager.send_other_message(data, send_user)
            else:
                data["username"] = user+'对全体'
                await ws_manager.broadcast_json(data)
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user, websocket)



from models.schemas import BaseUser
from fastapi import Header, HTTPException
from config import SECRET_KEY, ALGORITHM
from jose import JWTError, jwt


# 对token做了实际验证
async def get_current_user_token(websocket: WebSocket, token: Optional[str] = Header(...)) -> BaseUser:
    """token依赖请求头的token校验"""
    # 验证失败返回信息
    credentials_exc = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Authentication Failed",
        headers={"WWW-Authenticate": "Bearer"}
    )
    # 未登录返回信息
    not_login_exc = HTTPException(
        status_code=status.HTTP_403_FORBIDDEN,
        detail="NOT LOGIN OR TOKEN EXPIRED",
        headers={"test": "test"}
    )
    try:
        # 解析token
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username = payload.get("sub")
        if username is None:
            raise credentials_exc
        # 从redis获取用户token
        redis_token = await websocket.app.state.redis.get(username)
        if not redis_token or redis_token != token:
            raise not_login_exc
        return BaseUser(email=username)
    except JWTError:
        raise credentials_exc


# 对token做了实际验证， 放在Header中
@websocket_router.websocket("/ws2")
async def websocket_json_data2(websocket: WebSocket,  user: BaseUser = Depends(get_current_user_token)):
    await ws_manager.connect(user.email, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            send_user = data["username"]
            if send_user:
                # await ws_manager.send_other_message(data, send_user)
                data["username"] = '你'
                await ws_manager.send_personal_message_json(data, websocket)
                data["username"] = user.email
                await ws_manager.send_other_message(data, send_user)
            else:
                data["username"] = user.email+'对全体'
                await ws_manager.broadcast_json(data)
    except WebSocketDisconnect as e:
        await ws_manager.disconnect(user.email, websocket)
